import numpy as np
from langchain.embeddings import CacheBackedEmbeddings
from langchain.storage import LocalFileStore
from langchain_openai import OpenAIEmbeddings
import dotenv

dotenv.load_dotenv()

# 相似度计算函数
def cosine_similarty(vector1:list,vector2:list)->float:
    # 计算两个向量的点积
    dot_product = np.dot(vector1,vector2)
#     计算两个向量之间的长度
    vector1_norm = np.linalg.norm(vector1)
    vector2_norm = np.linalg.norm(vector2)
#     计算余弦相似度
    return dot_product /(vector1_norm * vector2_norm)

# 创建文本嵌入模型
embeddings = OpenAIEmbeddings(model = "text-embedding-3-small")
embedding_with_cache = CacheBackedEmbeddings.from_bytes_store(
    embeddings,
    LocalFileStore("./cache"),
    namespace=embeddings.model,
    query_embedding_cache=True
)

# 嵌入文本
query_vector = embedding_with_cache.embed_query("你好，我叫博小瑞，喜欢篮球？")
print(query_vector)
print(f"len = {len(query_vector)}")
# 嵌入文本列表
document_vector = embedding_with_cache.embed_documents(
    [
        "我叫博小瑞，喜欢打篮球",
        "这个喜欢打篮球的人叫，博小瑞",
        "乔布斯说过，求贤若渴、虚心若愚"
    ]
)
print(len(document_vector))